package org.skycloud.mybatis;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.AbstractList;
import java.util.List;
import java.util.Map;
import java.util.regex.Pattern;

import org.apache.ibatis.executor.result.DefaultResultHandler;
import org.apache.ibatis.logging.jdbc.ConnectionLogger;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.mapping.ParameterMode;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.ResultContext;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.apache.ibatis.type.JdbcType;
import org.apache.ibatis.type.TypeHandler;
import org.apache.ibatis.type.TypeHandlerRegistry;
import org.mybatis.spring.SqlSessionUtils;
import org.skycloud.mybatis.filter.Filter;
import org.skycloud.mybatis.pagination.PaginatedList;
import org.skycloud.mybatis.pagination.PaginatedListHandler;
import org.springframework.dao.support.PersistenceExceptionTranslator;

public final class SqlSessionTemplate extends org.mybatis.spring.SqlSessionTemplate {

	private static final String FROM = "from";
	private static final String ORDER = "order";
	private static final String COUNT_SUFFIX = "_count";
	private static final String SELECT_COUNT = "select count(*) ";
	private static final Pattern ORDER_BY_PATTERN = Pattern.compile("order\\p{javaWhitespace}+by");

	public SqlSessionTemplate(SqlSessionFactory sqlSessionFactory) {
		super(sqlSessionFactory);
	}

	public SqlSessionTemplate(SqlSessionFactory sqlSessionFactory, ExecutorType executorType) {
		super(sqlSessionFactory, executorType);
	}

	public SqlSessionTemplate(SqlSessionFactory sqlSessionFactory, ExecutorType executorType,
			PersistenceExceptionTranslator exceptionTranslator) {
		super(sqlSessionFactory, executorType, exceptionTranslator);
	}

	@Override
	public <K, V> Map<K, V> selectMap(String statement, String mapKey) {
		return selectMap(statement, null, mapKey);
	}

	@Override
	public <K, V> Map<K, V> selectMap(String statement, Object parameter, String mapKey) {
		return selectMap(statement, parameter, mapKey, RowBounds.DEFAULT);
	}

	@Override
	public <K, V> Map<K, V> selectMap(String statement, Object parameter, String mapKey, RowBounds rowBounds) {
		if (parameter instanceof Filter) {
			Filter filter = (Filter) parameter;
			rowBounds = new RowBounds(filter.getOffset(), filter.getLimit());
		} else if (rowBounds == null) {
			rowBounds = RowBounds.DEFAULT;
		}
		return super.selectMap(statement, parameter, mapKey, rowBounds);
	}

	@Override
	public <E> List<E> selectList(String statement) {
		return selectList(statement, null);
	}

	@Override
	public <E> List<E> selectList(String statement, Object parameter) {
		return selectList(statement, parameter, RowBounds.DEFAULT);
	}

	@Override
	public <T> List<T> selectList(String statement, Object parameter, RowBounds rowBounds) {
		if (parameter instanceof Filter) {
			Filter filter = (Filter) parameter;
			rowBounds = new RowBounds(filter.getOffset(), filter.getLimit());
		} else if (rowBounds == null) {
			rowBounds = RowBounds.DEFAULT;
		}

		return createPaginatedList(super.<T> selectList(statement, parameter, rowBounds), statement, parameter,
				rowBounds);
	}

	@Override
	public void select(String statement, ResultHandler handler) {
		select(statement, null, handler);
	}

	@Override
	public void select(String statement, Object parameter, ResultHandler handler) {
		select(statement, parameter, RowBounds.DEFAULT, handler);
	}

	@Override
	public void select(String statement, Object parameter, RowBounds rowBounds, ResultHandler handler) {
		if (parameter instanceof Filter) {
			Filter filter = (Filter) parameter;
			rowBounds = new RowBounds(filter.getOffset(), filter.getLimit());
		} else if (rowBounds == null) {
			rowBounds = RowBounds.DEFAULT;
		}

		if (handler instanceof PaginatedListHandler) {
			DefaultResultHandler defaultResultHandler = new DefaultResultHandler();
			super.select(statement, parameter, rowBounds, defaultResultHandler);

			final PaginatedList<?> paginatedList = createPaginatedList(defaultResultHandler.getResultList(), statement,
					parameter, rowBounds);

			handler.handleResult(new ResultContext() {

				@Override
				public Object getResultObject() {
					return paginatedList;
				}

				@Override
				public void stop() {
					throw new UnsupportedOperationException();
				}

				@Override
				public boolean isStopped() {
					throw new UnsupportedOperationException();
				}

				@Override
				public int getResultCount() {
					throw new UnsupportedOperationException();
				}
			});
		} else {
			super.select(statement, parameter, rowBounds, handler);
		}
	}

	private <T> PaginatedList<T> createPaginatedList(List<T> list, String statement, Object parameter,
			RowBounds rowBounds) {
		int totalSize;
		if (rowBounds.getOffset() == RowBounds.NO_ROW_OFFSET && rowBounds.getLimit() == RowBounds.NO_ROW_LIMIT) {
			totalSize = list.size();
		} else {
			boolean autoCount = false;
			if (parameter instanceof Filter) {
				autoCount = ((Filter) parameter).isAutoCount();
			}

			if (autoCount) {
				totalSize = count(statement, parameter);
			} else {
				totalSize = list.size();
			}
		}

		return new PaginatedListImpl<T>(list, rowBounds.getOffset(), rowBounds.getLimit(), totalSize);
	}

	/**
	 * Do sql count<br/>
	 * See {@code org.apache.ibatis.scripting.defaults.DefaultParameterHandler} for more details.
	 */
	@SuppressWarnings({ "unchecked", "rawtypes" })
	private int count(String statement, Object parameter) {
		// Get a SqlSession via SqlSessionUtils
		SqlSession sqlSession = SqlSessionUtils.getSqlSession(getSqlSessionFactory(), getExecutorType(),
				getPersistenceExceptionTranslator());

		Configuration configuration = getConfiguration();

		// If there's a count statement, count directly
		String countStatement = statement + COUNT_SUFFIX;
		if (configuration.getMappedStatementNames().contains(countStatement)) {
			return ((Number) selectOne(countStatement, parameter)).intValue();
		}

		// Create count sql automatically
		MappedStatement mappedStatement = configuration.getMappedStatement(statement);
		BoundSql boundSql = mappedStatement.getBoundSql(parameter);

		String sql = boundSql.getSql();
		String lowerCaseSql = sql.toLowerCase();

		int orderByIndex = lowerCaseSql.lastIndexOf(ORDER);
		boolean orderByFound = false;
		if (orderByIndex > 0 && ORDER_BY_PATTERN.matcher(lowerCaseSql).find(orderByIndex)) {
			orderByFound = true;
		}
		if (orderByFound) {
			sql = sql.substring(lowerCaseSql.indexOf(FROM), orderByIndex);
		} else {
			sql = sql.substring(lowerCaseSql.indexOf(FROM));
		}

		String countSql = SELECT_COUNT + sql;

		PreparedStatement pstmt = null;
		ResultSet rs = null;
		try {
			// This connection will be closed or returned to connection pool when SqlSession is closed, we do not close
			// this connection manually
			Connection conn = ConnectionLogger.newInstance(sqlSession.getConnection(),
					mappedStatement.getStatementLog());
			pstmt = conn.prepareStatement(countSql);

			Object parameterObject = boundSql.getParameterObject();
			TypeHandlerRegistry typeHandlerRegistry = configuration.getTypeHandlerRegistry();

			List<ParameterMapping> parameterMappings = boundSql.getParameterMappings();
			if (parameterMappings != null) {
				MetaObject metaObject = (parameterObject == null) ? null : configuration.newMetaObject(parameterObject);
				for (int i = 0; i < parameterMappings.size(); i++) {
					ParameterMapping parameterMapping = parameterMappings.get(i);
					if (parameterMapping.getMode() != ParameterMode.OUT) {
						Object value;
						String propertyName = parameterMapping.getProperty();
						if (boundSql.hasAdditionalParameter(propertyName)) {
							value = boundSql.getAdditionalParameter(propertyName);
						} else if (parameterObject == null) {
							value = null;
						} else if (typeHandlerRegistry.hasTypeHandler(parameterObject.getClass())) {
							value = parameterObject;
						} else {
							value = (metaObject == null) ? null : metaObject.getValue(propertyName);
						}
						TypeHandler typeHandler = parameterMapping.getTypeHandler();
						JdbcType jdbcType = parameterMapping.getJdbcType();
						if (value == null && jdbcType == null) {
							jdbcType = configuration.getJdbcTypeForNull();
						}
						typeHandler.setParameter(pstmt, i + 1, value, jdbcType);
					}
				}
			}

			rs = pstmt.executeQuery();
			return rs.next() ? rs.getInt(1) : 0;
		} catch (SQLException e) {
			throw new IllegalStateException("Can't get count with statement [" + statement
					+ "], you may surround your sql with [select * from (your sql) alias] or define a ["
					+ countStatement + "] statement in your MyBatis mapper xml file.", e);
		} finally {
			if (rs != null) {
				try {
					rs.close();
				} catch (SQLException ignore) {
				}
			}
			if (pstmt != null) {
				try {
					pstmt.close();
				} catch (SQLException ignore) {
				}
			}
			// Must close SqlSession here via SqlSessionUtils
			SqlSessionUtils.closeSqlSession(sqlSession, getSqlSessionFactory());
		}
	}

	@Override
	public <T> T getMapper(Class<T> type) {
		return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class<?>[] { type }, new MapperHandler<T>(
				type, super.getMapper(type), this)));
	}

	private static final class MapperHandler<T> implements InvocationHandler {

		private final Class<T> type;
		private final T mapper;
		private final SqlSession sqlSession;

		MapperHandler(Class<T> type, T mapper, SqlSession sqlSession) {
			this.type = type;
			this.mapper = mapper;
			this.sqlSession = sqlSession;
		}

		@Override
		public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
			Class<?> returnType = method.getReturnType();
			if (byte[].class.equals(returnType) || Byte[].class.equals(returnType)) {
				return sqlSession.selectOne(type.getName() + '.' + method.getName(), args);
			}

			return method.invoke(mapper, args);
		}
	}

	private static final class PaginatedListImpl<T> extends AbstractList<T> implements PaginatedList<T> {

		private static final long serialVersionUID = -6986381529211420399L;

		private final List<T> list;
		private final int offset;
		private final int limit;
		private final int totalSize;

		PaginatedListImpl(List<T> list, int offset, int limit, int totalSize) {
			this.list = list;
			this.offset = offset;
			this.limit = limit;
			this.totalSize = totalSize;
		}

		@Override
		public int getOffset() {
			return offset;
		}

		@Override
		public int getLimit() {
			return limit;
		}

		@Override
		public int getTotalSize() {
			return totalSize;
		}

		@Override
		public T get(int index) {
			return list.get(index);
		}

		@Override
		public int size() {
			return list.size();
		}
	}
}
